# PROJECT_NAME: vehicle_detection
# NAME: service_v6
# Author: XMX
# Date: 2025/11/11
# PROJECT_NAME: vehicle_detection
# NAME: service_v4_fixed
# Author: XMX
# Date: 2025/11/11
# smart_traffic_sandbox_service_fixed.py
from fastapi import FastAPI, File, UploadFile, HTTPException, BackgroundTasks, WebSocket, WebSocketDisconnect
from fastapi.responses import StreamingResponse, JSONResponse
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
import cv2
import numpy as np
import base64
import json
import uuid
import asyncio
import time
import os
from typing import Dict, List, Optional, Tuple
from datetime import datetime
from ultralytics import YOLO
import torch
import logging
from collections import defaultdict, deque

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

app = FastAPI(
    title="智慧交通检测服务",
    description="沙盘模拟验证系统 - 车辆识别与分级告警",
    version="5.0"
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 创建必要的目录
os.makedirs("temp_videos", exist_ok=True)
os.makedirs("output_videos", exist_ok=True)


class VehicleIdentity:
    """车辆身份信息"""

    def __init__(self, vehicle_id: str, vehicle_type: str):
        self.id = vehicle_id
        self.type = vehicle_type
        self.first_seen = datetime.now()
        self.last_seen = datetime.now()
        self.track_history = []  # 轨迹历史


class AlertInfo:
    """告警信息"""

    def __init__(self, alert_id: str, alert_type: str, bbox: List[float], confidence: float):
        self.id = alert_id
        self.type = alert_type  # 'accident'
        self.bbox = bbox
        self.confidence = confidence
        self.start_time = datetime.now()
        self.last_update = datetime.now()
        self.continuous_frames = 1
        self.confirmed = False  # 是否确认为真实警报


class SandboxTrafficDetector:
    def __init__(self, model_path: str):
        self.model = None

        # 9种车辆类型
        self.vehicle_types = [
            'yellow_car', 'white_car', 'green_car', 'red_car',
            'bus', 'truck', 'pink_car', 'blue_car', 'gray_car'
        ]

        # 车辆显示颜色映射
        self.vehicle_colors = {
            'yellow_car': (0, 255, 255),  # 黄色
            'white_car': (255, 255, 255),  # 白色
            'green_car': (0, 255, 0),  # 绿色
            'red_car': (0, 0, 255),  # 红色
            'bus': (255, 0, 0),  # 蓝色
            'truck': (128, 0, 128),  # 紫色
            'pink_car': (255, 0, 255),  # 粉色
            'blue_car': (255, 0, 0),  # 蓝色
            'gray_car': (128, 128, 128)  # 灰色
        }

        self.device = 'cuda' if torch.cuda.is_available() else 'cpu'

        # 实时状态
        self.current_vehicles = {}  # vehicle_id -> VehicleIdentity
        self.active_alerts = {}  # alert_id -> AlertInfo

        # 告警配置 - 只保留碰撞检测
        self.alert_config = {
            'accident': {
                'confidence_threshold': 0.7,
                'confirm_frames': 5,
            }
        }

        # 车辆跟踪优化
        self.tracking_threshold = 0.3
        self.max_track_points = 30

        self.load_model(model_path)
        logger.info(f"设备: {self.device}")

    def load_model(self, model_path: str):
        try:
            self.model = YOLO(model_path)
            if self.device == 'cuda':
                self.model.to('cuda')
            logger.info("YOLO模型加载成功")
        except Exception as e:
            logger.error(f"模型加载失败: {e}")
            raise

    def detect_vehicles(self, image: np.ndarray, conf_threshold: float = 0.25) -> List[Dict]:
        """检测图像中的车辆"""
        try:
            results = self.model(image, conf=conf_threshold, device=self.device, verbose=False)

            detections = []
            if results and len(results) > 0 and results[0].boxes is not None:
                for box in results[0].boxes:
                    cls_id = int(box.cls[0])
                    # 直接使用类别ID对应的车辆类型
                    vehicle_type = self.vehicle_types[cls_id] if cls_id < len(self.vehicle_types) else 'unknown'
                    confidence = float(box.conf[0])
                    bbox = box.xyxy[0].cpu().numpy().tolist()

                    # 生成车辆唯一ID (基于位置和车辆类型)
                    center_x = (bbox[0] + bbox[2]) / 2
                    center_y = (bbox[1] + bbox[3]) / 2
                    vehicle_id = f"{vehicle_type}_{int(center_x // 10)}_{int(center_y // 10)}"

                    detection = {
                        'vehicle_id': vehicle_id,
                        'vehicle_type': vehicle_type,
                        'confidence': confidence,
                        'bbox': bbox,
                        'center': {'x': center_x, 'y': center_y}
                    }
                    detections.append(detection)

            return detections
        except Exception as e:
            logger.error(f"车辆检测失败: {e}")
            return []

    def detect_accidents(self, detections: List[Dict], image_shape: Tuple[int, int]) -> List[AlertInfo]:
        """检测交通事故"""
        accidents = []

        # 简化的事故检测逻辑：车辆重叠或异常接近
        for i, det1 in enumerate(detections):
            for det2 in detections[i + 1:]:
                iou = self.calculate_iou(det1['bbox'], det2['bbox'])

                # 如果车辆重叠区域较大，认为是事故
                if iou > 0.15:
                    # 计算事故区域
                    bbox1 = det1['bbox']
                    bbox2 = det2['bbox']
                    accident_bbox = [
                        min(bbox1[0], bbox2[0]),
                        min(bbox1[1], bbox2[1]),
                        max(bbox1[2], bbox2[2]),
                        max(bbox1[3], bbox2[3])
                    ]

                    # 置信度基于重叠程度
                    confidence = min(iou * 3, 1.0)

                    if confidence > self.alert_config['accident']['confidence_threshold']:
                        alert_id = f"accident_{int(accident_bbox[0])}_{int(accident_bbox[1])}"
                        accidents.append(AlertInfo(alert_id, 'accident', accident_bbox, confidence))

        return accidents

    def calculate_iou(self, bbox1: List[float], bbox2: List[float]) -> float:
        """计算IOU"""
        x1 = max(bbox1[0], bbox2[0])
        y1 = max(bbox1[1], bbox2[1])
        x2 = min(bbox1[2], bbox2[2])
        y2 = min(bbox1[3], bbox2[3])

        if x2 < x1 or y2 < y1:
            return 0.0

        intersection = (x2 - x1) * (y2 - y1)
        area1 = (bbox1[2] - bbox1[0]) * (bbox1[3] - bbox1[1])
        area2 = (bbox2[2] - bbox2[0]) * (bbox2[3] - bbox2[1])

        union = area1 + area2 - intersection
        return intersection / union if union > 0 else 0.0

    def update_vehicle_tracking(self, detections: List[Dict]):
        """更新车辆跟踪"""
        current_time = datetime.now()
        updated_vehicles = {}

        for detection in detections:
            vehicle_id = detection['vehicle_id']
            vehicle_type = detection['vehicle_type']

            if vehicle_id in self.current_vehicles:
                # 更新现有车辆
                vehicle = self.current_vehicles[vehicle_id]
                vehicle.last_seen = current_time
                vehicle.track_history.append(detection['center'])
                # 保持轨迹长度合理
                if len(vehicle.track_history) > self.max_track_points:
                    vehicle.track_history.pop(0)
            else:
                # 新车辆
                vehicle = VehicleIdentity(vehicle_id, vehicle_type)
                vehicle.track_history.append(detection['center'])

            updated_vehicles[vehicle_id] = vehicle

        self.current_vehicles = updated_vehicles

    def update_alerts(self, new_accidents: List[AlertInfo]):
        """更新告警状态 - 只处理碰撞告警"""
        current_time = datetime.now()
        updated_alerts = {}

        # 只处理当前帧的告警
        for alert in new_accidents:
            if alert.id in self.active_alerts:
                # 更新现有告警
                existing_alert = self.active_alerts[alert.id]
                existing_alert.last_update = current_time
                existing_alert.continuous_frames += 1

                # 检查是否确认为真实警报
                required_frames = self.alert_config['accident']['confirm_frames']
                if existing_alert.continuous_frames >= required_frames:
                    existing_alert.confirmed = True

                updated_alerts[alert.id] = existing_alert
            else:
                # 新告警
                alert.start_time = current_time
                alert.last_update = current_time
                updated_alerts[alert.id] = alert

        self.active_alerts = updated_alerts

    def draw_detection_results(self, image: np.ndarray, detections: List[Dict],
                               frame_count: int, fps: float) -> np.ndarray:
        """绘制检测结果到图像"""
        result_image = image.copy()
        height, width = result_image.shape[:2]

        # 绘制车辆检测框
        for detection in detections:
            bbox = detection['bbox']
            vehicle_type = detection['vehicle_type']
            confidence = detection['confidence']

            # 获取车辆颜色
            color = self.vehicle_colors.get(vehicle_type, (0, 255, 0))

            # 绘制边界框
            cv2.rectangle(result_image,
                          (int(bbox[0]), int(bbox[1])),
                          (int(bbox[2]), int(bbox[3])),
                          color, 2)

            # 使用英文标签避免中文乱码问题
            label = f"{vehicle_type} {confidence:.2f}"
            label_size = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 2)[0]

            cv2.rectangle(result_image,
                          (int(bbox[0]), int(bbox[1]) - label_size[1] - 10),
                          (int(bbox[0]) + label_size[0], int(bbox[1])),
                          color, -1)

            cv2.putText(result_image, label,
                        (int(bbox[0]), int(bbox[1]) - 5),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 2)

        # 绘制告警框 - 只在当前帧有告警时显示
        for alert_id, alert in self.active_alerts.items():
            bbox = alert.bbox
            alert_color = (0, 0, 255) if alert.confirmed else (0, 165, 255)

            # 绘制告警框
            cv2.rectangle(result_image,
                          (int(bbox[0]), int(bbox[1])),
                          (int(bbox[2]), int(bbox[3])),
                          alert_color, 4)

            # 告警标签
            status_text = "CONFIRMED" if alert.confirmed else "DETECTING"
            alert_text = f"ACCIDENT[{status_text}]"
            alert_size = cv2.getTextSize(alert_text, cv2.FONT_HERSHEY_SIMPLEX, 0.7, 2)[0]

            cv2.rectangle(result_image,
                          (int(bbox[0]), int(bbox[1]) - alert_size[1] - 10),
                          (int(bbox[0]) + alert_size[0], int(bbox[1])),
                          alert_color, -1)

            cv2.putText(result_image, alert_text,
                        (int(bbox[0]), int(bbox[1]) - 5),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2)

        # 绘制统计信息面板
        info_panel = np.zeros((200, 400, 3), dtype=np.uint8)

        info_lines = [
            f"Frame: {frame_count}",
            f"FPS: {fps:.1f}",
            f"Vehicles: {len(self.current_vehicles)}",
            f"Alerts: {len([a for a in self.active_alerts.values() if a.confirmed])}",
            f"Device: {self.device}",
            "---Vehicle Distribution---"
        ]

        # 添加车辆类型统计
        type_count = defaultdict(int)
        for vehicle in self.current_vehicles.values():
            type_count[vehicle.type] += 1

        for vehicle_type in self.vehicle_types:
            count = type_count.get(vehicle_type, 0)
            if count > 0:
                info_lines.append(f"{vehicle_type}: {count}")

        # 添加告警统计
        if self.active_alerts:
            info_lines.append("---Alert Statistics---")
            alert_count = len([a for a in self.active_alerts.values() if a.confirmed])
            info_lines.append(f"Accidents: {alert_count}")

        for i, line in enumerate(info_lines):
            color = (0, 255, 0) if i < 5 else (255, 255, 255)
            cv2.putText(info_panel, line, (10, 20 + i * 20),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 1)

        # 将信息面板叠加到主图像
        result_image[10:210, 10:410] = info_panel

        return result_image

    def get_realtime_statistics(self) -> Dict:
        """获取实时统计信息"""
        vehicle_type_count = defaultdict(int)
        for vehicle in self.current_vehicles.values():
            vehicle_type_count[vehicle.type] += 1

        confirmed_alerts = [alert for alert in self.active_alerts.values() if alert.confirmed]

        return {
            'timestamp': datetime.now().isoformat(),
            'vehicle_count': len(self.current_vehicles),
            'vehicle_distribution': dict(vehicle_type_count),
            'active_alerts': len(confirmed_alerts),
            'alert_details': [
                {
                    'type': alert.type,
                    'bbox': alert.bbox,
                    'confidence': alert.confidence,
                    'duration': (datetime.now() - alert.start_time).total_seconds(),
                    'confirmed': alert.confirmed
                }
                for alert in confirmed_alerts
            ]
        }


# 全局检测器实例
detector = SandboxTrafficDetector(
    model_path=r'F:\MyProject\Python\vehicle_detection\runs\detect\train2\weights\best.pt'
)

# 任务存储和实时数据
analysis_tasks = {}
realtime_stats = {
    'last_update': None,
    'statistics': None
}


@app.on_event("startup")
async def startup_event():
    logger.info("服务初始化完成")


@app.get("/")
async def root():
    return {
        "service": "智慧交通检测服务",
        "version": "5.0",
        "status": "运行中",
        "device": detector.device,
        "supported_vehicles": detector.vehicle_types
    }


@app.get("/health")
async def health_check():
    return {
        "status": "healthy",
        "timestamp": datetime.now().isoformat(),
        "model_loaded": detector.model is not None,
        "device": detector.device,
        "supported_vehicles": len(detector.vehicle_types)
    }


@app.post("/api/analyze/video")
async def analyze_video_api(
        background_tasks: BackgroundTasks,
        file: UploadFile = File(...),
        conf_threshold: float = 0.25
):
    """视频分析API"""
    try:
        task_id = str(uuid.uuid4())

        # 保存上传的视频
        video_path = os.path.join("temp_videos", f"{task_id}_{file.filename}")
        with open(video_path, "wb") as f:
            content = await file.read()
            f.write(content)

        # 设置输出路径
        output_filename = f"processed_{task_id}.mp4"
        output_path = os.path.join("output_videos", output_filename)

        # 初始化任务状态
        analysis_tasks[task_id] = {
            "status": "processing",
            "video_path": video_path,
            "output_path": output_path,
            "start_time": datetime.now().isoformat()
        }

        # 在后台处理视频
        background_tasks.add_task(process_video_task, task_id, video_path, output_path, conf_threshold)

        return {
            "success": True,
            "task_id": task_id,
            "message": "视频分析任务已开始",
            "status_url": f"/api/tasks/{task_id}",
            "output_url": f"/api/video/output/{task_id}"
        }

    except Exception as e:
        logger.error(f"视频分析失败: {e}")
        raise HTTPException(status_code=500, detail=f"视频分析失败: {str(e)}")


async def process_video_task(task_id: str, video_path: str, output_path: str, conf_threshold: float):
    """处理视频任务"""
    try:
        cap = cv2.VideoCapture(video_path)
        if not cap.isOpened():
            raise Exception("无法打开视频文件")

        # 获取视频信息
        fps = cap.get(cv2.CAP_PROP_FPS)
        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

        # 设置输出视频
        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
        out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))

        frame_count = 0
        fps_calculator = FPSCalculator()
        last_stat_update = time.time()

        while True:
            start_time = time.time()

            ret, frame = cap.read()
            if not ret:
                break

            frame_count += 1

            # 执行车辆检测
            detections = detector.detect_vehicles(frame, conf_threshold)

            # 更新车辆跟踪
            detector.update_vehicle_tracking(detections)

            # 检测告警 - 只检测碰撞
            accidents = detector.detect_accidents(detections, frame.shape)

            # 更新告警状态
            detector.update_alerts(accidents)

            # 每秒更新一次统计信息
            current_time = time.time()
            if current_time - last_stat_update >= 1.0:
                realtime_stats['statistics'] = detector.get_realtime_statistics()
                realtime_stats['last_update'] = datetime.now().isoformat()
                last_stat_update = current_time

            # 绘制检测结果
            current_fps = fps_calculator.update()
            result_frame = detector.draw_detection_results(frame, detections, frame_count, current_fps)

            # 写入输出视频
            out.write(result_frame)

            # 控制处理速度
            processing_time = time.time() - start_time
            wait_time = max(0, 1 / fps - processing_time)
            if wait_time > 0:
                time.sleep(wait_time)

        # 完成处理
        cap.release()
        out.release()

        analysis_tasks[task_id].update({
            "status": "completed",
            "output_path": output_path,
            "completion_time": datetime.now().isoformat(),
            "total_frames": frame_count
        })

        # 清理临时文件
        if os.path.exists(video_path):
            os.remove(video_path)

        logger.info(f"视频处理完成: {task_id}, 总帧数: {frame_count}")

    except Exception as e:
        analysis_tasks[task_id].update({
            "status": "failed",
            "error": str(e),
            "completion_time": datetime.now().isoformat()
        })
        logger.error(f"视频处理失败: {e}")


@app.get("/api/tasks/{task_id}")
async def get_task_status(task_id: str):
    """获取任务状态"""
    if task_id not in analysis_tasks:
        raise HTTPException(status_code=404, detail="任务不存在")

    task_info = analysis_tasks[task_id]
    response = {
        "task_id": task_id,
        "status": task_info["status"],
        "start_time": task_info["start_time"]
    }

    if task_info["status"] == "completed":
        response["output_path"] = task_info["output_path"]
        response["completion_time"] = task_info["completion_time"]
        response["total_frames"] = task_info["total_frames"]
        response["download_url"] = f"/api/video/output/{task_id}"
    elif task_info["status"] == "failed":
        response["error"] = task_info["error"]
        response["completion_time"] = task_info["completion_time"]

    return response


@app.get("/api/video/output/{task_id}")
async def get_processed_video(task_id: str):
    """获取处理后的视频文件"""
    if task_id not in analysis_tasks or analysis_tasks[task_id]["status"] != "completed":
        raise HTTPException(status_code=404, detail="视频未就绪")

    output_path = analysis_tasks[task_id]["output_path"]
    if not os.path.exists(output_path):
        raise HTTPException(status_code=404, detail="视频文件不存在")

    def iterfile():
        with open(output_path, "rb") as f:
            yield from f

    return StreamingResponse(iterfile(), media_type="video/mp4")


@app.get("/api/realtime/stats")
async def get_realtime_statistics():
    """获取实时统计信息"""
    if realtime_stats['statistics'] is None:
        return {
            "timestamp": datetime.now().isoformat(),
            "vehicle_count": 0,
            "vehicle_distribution": {},
            "active_alerts": 0,
            "alert_details": []
        }

    return realtime_stats['statistics']


class FPSCalculator:
    def __init__(self):
        self.start_time = time.time()
        self.frame_count = 0

    def update(self):
        self.frame_count += 1
        current_time = time.time()
        elapsed = current_time - self.start_time

        if elapsed > 0:
            fps = self.frame_count / elapsed
            if elapsed >= 1.0:
                self.start_time = current_time
                self.frame_count = 0
            return fps
        return 0


if __name__ == "__main__":
    import uvicorn

    print("=" * 60)
    print(" 接口文档: GET http://localhost:8000/docs")
    print(" 实时统计: GET http://localhost:8000/api/realtime/stats")
    print(" 视频上传: POST http://localhost:8000/api/analyze/video")
    print(" 任务状态: GET http://localhost:8000/api/tasks/{task_id}")
    print(" 视频下载: GET http://localhost:8000/api/video/output/{task_id}")
    print(" 健康检查: GET http://localhost:8000/health")
    print("=" * 60)

    uvicorn.run(app, host="0.0.0.0", port=8000, reload=False)